#include <limits>
#include <vector>

#include <glog/logging.h>

#include <QtGlobal>

#include <botan/dh.h>
#include <botan/pubkey.h>


#include "octopod.h"
#include "securechannel.h"
#include "securechannel_auth_message.pb.h"
#include "crypto.h"
#include "share.h"
#include "peer.h"
#include "store.h"

namespace octopod {

SecureChannel::SecureChannel(QSharedPointer<Socket> sock, QSharedPointer<Peer> peer,
                             QSharedPointer<Share> share,
                             const Botan::SymmetricKey &session_key) 
    : sock_(sock),
      peer_(peer),
      share_(share),
      session_key_(session_key),
      outgoing_message_counter_(0),
      incoming_message_counter_(0) {
}


SecureChannel *SecureChannel::createOutgoing(Socket *s, QSharedPointer<Peer> peer,
                                             QSharedPointer<Share> share) {
  CHECK_NOTNULL(s);
  CHECK(!peer.isNull());
  CHECK(!share.isNull());

  // Wrap the socket in a shared pointer so that it will be destroyed for sure
  // regardless of whether authentication succeeds or not.
  QSharedPointer<Socket> sock(s);

  // If authentication succeeds...
  Botan::SymmetricKey session_key;
  if (authenticateOutgoing(sock, peer, share, &session_key)) {
    // create a secure channel and return it
    return new SecureChannel(sock, peer, share, session_key);
  }

  // Otherwise return NULL and allow sock to be destroyed.
  return NULL;
}

SecureChannel *SecureChannel::createIncoming(Socket *s) {
  CHECK_NOTNULL(s);

  // Wrap the socket to ensure its destruction, like in createOutgoing.
  QSharedPointer<Socket> sock(s);

  // Since we don't know the peer or share yet, authentication will have
  // to figure that out and return it to us.  There is no good way to return
  // two QSharedPointers from a method, so we'll have to use a pair to do it.
  // (A) NULL shared pointer(s) returned from authentication indicate(s) that
  // authentication failed.
  Botan::SymmetricKey session_key;
  std::pair< QSharedPointer<Peer>, QSharedPointer<Share> > ps =
       authenticateIncoming(sock, &session_key);

  QSharedPointer<Peer> peer = ps.first;
  QSharedPointer<Share> share = ps.second;

  // If authentication succeeds...
  if (!peer.isNull() && !share.isNull()) {
    return new SecureChannel(sock, peer, share, session_key);
  }

  // Otherwise return NULL and allow sock to be destroyed.
  return NULL;
}

// Sends a message via the secure channel.
// Returns true upon success.
bool SecureChannel::send(const std::string &msg) {
  std::string tmpmsg = msg;
  
  // handle the message counter in network byte order
  if (outgoing_message_counter_ >= UINT32_MAX) {
    DLOG(WARNING) << Q_FUNC_INFO << "too many messages sent on this channel.";
    return false;
  }

  uint32_t ctr = htonl(outgoing_message_counter_);
  outgoing_message_counter_++;
  tmpmsg.append((const char *) &ctr, sizeof(ctr));
 
  if (!sendEncryptedMessage(sock_, session_key_, tmpmsg)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't send message";
    return false;
  }

  return true;
}

// Receives a message via the secure channel.
// Stores the message in the given string.
bool SecureChannel::recv(std::string *msg) {
  if (!recvEncryptedMessage(sock_, session_key_, msg)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't receive message";
    return false;
  }

  // ensure the message at least contains a messgae counter
  if (msg->size() < sizeof(uint32_t)) {
    DLOG(WARNING) << Q_FUNC_INFO << "got an irregularly small message";
    return false;
  }

  // position of the message counter
  size_t pos = msg->size() - sizeof(uint32_t);
  uint32_t ctr = *((uint32_t *) (msg->data() + pos));
  ctr = ntohl(ctr);
 
  // check that message counter is what we expect next
  if (ctr != incoming_message_counter_) {
    DLOG(WARNING) << Q_FUNC_INFO << "got a message with an unexpected message counter, rejecting";
    return false;
  }

  // remove the message counter
  msg->resize(msg->size() - sizeof(uint32_t));

  return true;
}

bool SecureChannel::sendEncryptedMessage(QSharedPointer<Socket> sock,
                                         const Botan::SymmetricKey &key,
                                         const std::string &data) {
  CHECK(!sock.isNull());
  CHECK_GT(data.size(), 0);
  CHECK_LE(data.size() + crypto::kAEOverhead + crypto::kZlibOverhead, INT32_MAX);

  std::string ciphertext, init_vector;

  // encrypt the message
  if (!crypto::authEncrypt(key, data, &ciphertext, &init_vector)) {
    DLOG(WARNING) << Q_FUNC_INFO << "authenticated encryption failed.";
    return false;
  }
  
  // compute the message's length in network byte order
  uint32_t msglen = ciphertext.size();
  msglen = htonl(msglen);
  
  std::string msglen_str((const char *) &msglen, sizeof(msglen));

  int err;
  // send the message length, encrypted message IV, and ciphertext
  if (((err = sock->send(msglen_str)) < 0) || err != (int) msglen_str.size()) {
    DLOG(WARNING) << Q_FUNC_INFO << "sending the IV for the message failed.";
    return false;
  }

  if (((err = sock->send(init_vector)) < 0) || err != (int) init_vector.size()) {
    DLOG(WARNING) << Q_FUNC_INFO << "sending the IV for the message failed.";
    return false;
  }

  if (((err = sock->send(ciphertext)) < 0) || err != (int) ciphertext.size()) {
    DLOG(WARNING) << Q_FUNC_INFO << "sending the ciphertext for the message failed.";
    return false;
  }

  return true;
}

bool SecureChannel::recvEncryptedMessageData(QSharedPointer<Socket> sock,
                                             std::string *init_vector,
                                             std::string *ciphertext) {
  // receive message length
  uint32_t msglen;
  std::string msglen_str;
  int err;

  if (((err = sock->recv(sizeof(msglen), &msglen_str)) < 0) || err != (int) sizeof(msglen)) {
    DLOG(WARNING) << Q_FUNC_INFO << "reading the message length failed.";
    return false;
  }

  msglen = *((uint32_t *) msglen_str.data());
  msglen = ntohl(msglen);

  if (msglen > INT32_MAX) {
    DLOG(WARNING) << Q_FUNC_INFO << "received a ciphertext that was too large";
    return false;
  }

  // receive message IV
  if (sock->recv(crypto::kIVLength, init_vector) != crypto::kIVLength) {
    DLOG(WARNING) << Q_FUNC_INFO << "reading the IV for the message failed.";
    return false;
  }

  // receive message ciphertext
  if (((err = sock->recv(msglen, ciphertext)) < 0) || err != (int) msglen) {
    DLOG(WARNING) << Q_FUNC_INFO << "reading the ciphertext for the message failed.";
    return false;
  }

  return true;
}

bool SecureChannel::recvEncryptedMessage(QSharedPointer<Socket> sock,
                                         const Botan::SymmetricKey &key,
                                         std::string *data) {
  CHECK(!sock.isNull());

  std::string ciphertext, init_vector;

  if (!recvEncryptedMessageData(sock, &init_vector, &ciphertext)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't receive message data.";
    return false;
  }

  // decrypt and process the message
  if (!crypto::authDecrypt(key, ciphertext, init_vector, data)) {
    DLOG(WARNING) << Q_FUNC_INFO << "decrypting the message failed.";
    return false;
  }

  return true;
}

QSharedPointer<Share> SecureChannel::recvEncryptedMessage(QSharedPointer<Socket> sock, 
                                                          std::string *data) {
  CHECK(!sock.isNull());

  QSharedPointer<Share> nullptr(NULL); // to return upon error

  std::string ciphertext, init_vector;

  if (!recvEncryptedMessageData(sock, &init_vector, &ciphertext)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't receive message data.";
    return nullptr;
  }

  QSharedPointer<Share> share; // the right share, unknown at the moment

  // we need to loop to figure out which share key to use
  const std::vector< QSharedPointer<Share> > &all_shares = store::getAllShares();

  for (std::vector< QSharedPointer<Share> >::const_iterator i = all_shares.begin(); i != all_shares.end(); ++i) {
    QSharedPointer<Share> test_share = *i;
    const Botan::SymmetricKey &test_key = test_share->peer_auth_key();

    if (crypto::authDecrypt(test_key, ciphertext, init_vector, data)) {
      // found the right share, so we're done.
      share = test_share;
      break;
    }
  }

  if (share.isNull()) {
    DLOG(WARNING) << Q_FUNC_INFO << "unable to find a share key to decrypt incoming messages.";
    return nullptr;
  }

  return share;
}

// Authenticates an already connected socket with the given peer and share.
bool SecureChannel::authenticateOutgoing(QSharedPointer<Socket> sock,
                                         QSharedPointer<Peer> peer,
                                         QSharedPointer<Share> share,
                                         Botan::SymmetricKey *session_key) {
  CHECK(!sock.isNull());

  QSharedPointer<Peer> me = share->me();
  const Botan::SymmetricKey &peer_auth_key = share->peer_auth_key();

  // Follows steps as labeled in secure channel authentication protocol.
  // In that protocol, we're host A since A is the outgoing host.
  //
  // 1. A: generates random nonce Nonce_A, computes H(Nonce_A)
  std::string nonce_a;
  if (!crypto::randomGen(kNonceLength, &nonce_a)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't generate random nonce for authentication at step 1";
    return false;
  }

  // compute hash of nonce
  std::string nonce_a_hash = crypto::hash(nonce_a);

  // generate an empty nonce string
  std::string empty_nonce(kNonceLength, '\0');

  // 2. A->B: C_1 = AE_{PA-key}(H(A_pk), H(B_pk), H(share ID), H(Nonce_A), 0)
  // set up message and send
  SecureChannelAuthMessage outmsg;
  setAuthMsg(me->keyhash(), peer->keyhash(), share->idhash(), nonce_a_hash,
             empty_nonce, &outmsg);
  
  if (!sendEncryptedMessage(sock, peer_auth_key, outmsg.SerializeAsString())) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't send at auth step 2.";
    return false;
  }

  // 3. B: decrypts C_1 using all available PA-key values (for all shares B is a member of) 
  // (do nothing)

  // 4. B: stores H(Nonce_A) and the other values
  // (do nothing)

  // 5. B: generates random nonce Nonce_B, computes H(Nonce_B)
  // (do nothing)

  // 6. B->A: C_2 = AE_{PA-key}(H(B_pk), H(A_pk), H(share ID), H(Nonce_A), H(Nonce_B))
  // receive to inmsg
  std::string instr;
  
  if (!recvEncryptedMessage(sock, peer_auth_key, &instr)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't receive at auth step 6.";
    return false;
  }

  // 7. A: checks the same conditions as B did for C_1
  // process the received message and ensure it's in the right format.
  SecureChannelAuthMessage inmsg;
  if (!processAuthMsg(instr, me, share, &inmsg)) {
    DLOG(WARNING) << Q_FUNC_INFO << "auth message was in a bad format at auth step 7.";
    return false;
  }

  // 8. A: stores H(Nonce_B) and the other values
  // keep inmsg around.
  std::string nonce_b_hash = inmsg.second_nonce();

  // 9. A->B: C_3 = AE_{PA-key}(H(A_pk), H(B_pk), H(share ID), Nonce_A, H(Nonce_B))
  // set up message and send
  SecureChannelAuthMessage outmsg2;
  setAuthMsg(me->keyhash(), peer->keyhash(), share->idhash(), nonce_a,
             nonce_b_hash, &outmsg2);
  
  if (!sendEncryptedMessage(sock, peer_auth_key, outmsg2.SerializeAsString())) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't send at auth step 9.";
    return false;
  }

  // 10. B: checks all the above conditions and checks that H(Nonce_A) == previously stored H(Nonce_A)
  // (do nothing)

  // 11. B->A: C_4 = AE_{PA-key}(H(B_pk), H(A_pk), H(share ID), Nonce_A, Nonce_B)
  // receive to inmsg2
  if (!recvEncryptedMessage(sock, peer_auth_key, &instr)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't receive at auth step 11.";
    return false;
  }
  
  // 12. A: checks all the above conditions and checks that H(Nonce_B) == previously stored H(Nonce_B)
  SecureChannelAuthMessage inmsg2;
  if (!processAuthMsg(instr, me, share, &inmsg2)) {
    DLOG(WARNING) << Q_FUNC_INFO << "auth message was in a bad format at auth step 12.";
    return false;
  }

  if (crypto::hash(inmsg2.second_nonce()) != inmsg.second_nonce() ||
      inmsg2.first_nonce() != outmsg2.first_nonce()) {
    DLOG(WARNING) << Q_FUNC_INFO << "nonce failed check at auth step 12.";
    return false;
  }

  // save the final authentication message to use as the salt for the KDF.
  const std::string salt = instr;

  // 13. A: generates Diffie-Hellman key exchange parameters
  // set up DH group and private/public values
  // normally this would be done in crypto.cc but we never use DH elsewhere.
  // could be moved later.
  Botan::AutoSeeded_RNG rng;
  
  Botan::DL_Group shared_domain(crypto::kDHDomain);
  Botan::DH_PrivateKey private_a(rng, shared_domain);
  Botan::MemoryVector<Botan::byte> public_a = private_a.public_value();

  // 14. A->B: AE_{PA-key}(PKE_{B_pk}(public_a))
  std::string public_a_str(reinterpret_cast<const char *>(public_a.begin()), public_a.size());

  std::string enc_public_a_str; // the public_a value encrypted under B's pubkey
  if (!crypto::pubkeyEncrypt(peer->pubkey(), public_a_str, &enc_public_a_str)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't encrypt at auth step 14.";
    return false;
  }
  
  // encrypt and send PKE_{B_pk}(public_a)
  if (!sendEncryptedMessage(sock, peer_auth_key, enc_public_a_str)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't send at auth step 14.";
    return false;
  }

  // 15. B: generates Diffie-Hellman key exchange parameters
  // (do nothing)

  // 16. B->A: AE_{PA-key}(PKE_{A_pk}(public_b))
  // receive the message and produce the corresponding MemoryVector
  if (!recvEncryptedMessage(sock, peer_auth_key, &instr)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't receive at auth step 16.";
    return false;
  }

  // decrypt using our private key.
  std::string public_b_str;
  if (!crypto::pubkeyDecrypt(me->privkey(), instr, &public_b_str)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't decrypt at auth step 16.";
    return false;
  }

  // 17. Generate session key from DH and salt
  Botan::MemoryVector<Botan::byte> public_b(reinterpret_cast<const Botan::byte *>(public_b_str.data()), public_b_str.size());
  Botan::PK_Key_Agreement ka(private_a, crypto::kDH_KDF);

  try {
    *session_key = ka.derive_key(crypto::kKeyLength, public_b, salt);
  } catch (Botan::Exception &e) {
    DLOG(ERROR) << Q_FUNC_INFO << "Botan exception: " << e.what();
    return false;
  }
  
  return true;
}

// Since for an incoming connection we don't know about the peer or share
// we need to figure it out and then return the peer and share the
// connection pertains to.  If we can't find any, then we return NULL
// for both and the incoming connection can't be authenticated.
std::pair< QSharedPointer<Peer>, QSharedPointer<Share> > SecureChannel::authenticateIncoming(QSharedPointer<Socket> sock, Botan::SymmetricKey *session_key) {
  CHECK(!sock.isNull());

  // create a NULL pointer pair to return upon error
  std::pair< QSharedPointer<Peer>, QSharedPointer<Share> > fail;

  // Follows steps as labeled in secure channel authentication protocol.
  // In that protocol, we're host B since A is the initiating host.
  //
  // 1. A: generates random nonce Nonce_A, computes H(Nonce_A)
  // (do nothing)

  // 2. A->B: C_1 = AE_{PA-key}(H(A_pk), H(B_pk), H(share ID), H(Nonce_A), 0)
  // receive message
  std::string instr;
  QSharedPointer<Share> share;

  // 3. B: decrypts C_1 using all available PA-key values (for all shares B is a member of) 
  share = recvEncryptedMessage(sock, &instr);
  if (!share.isNull()) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't receive at auth step 2.";
    return fail;
  }
  
  // 4. B: stores H(Nonce_A) and the other values
  QSharedPointer<Peer> me = share->me();
  SecureChannelAuthMessage inmsg;
  if (!processAuthMsg(instr, me, share, &inmsg)) {
    DLOG(WARNING) << Q_FUNC_INFO << "auth message was in a bad format at auth step 2.";
    return fail;
  }
  std::string nonce_a_hash = inmsg.first_nonce();

  QSharedPointer<Peer> peer = store::getPeerByHash(inmsg.from_user_keyhash());

  const Botan::SymmetricKey peer_auth_key = share->peer_auth_key();

  // 5. B: generates random nonce Nonce_B, computes H(Nonce_B)
  std::string nonce_b;
  if (!crypto::randomGen(kNonceLength, &nonce_b)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't generate random nonce for authentication at step 5";
    return fail;
  }

  // compute hash of nonce
  std::string nonce_b_hash = crypto::hash(nonce_b);

  // 6. B->A: C_2 = AE_{PA-key}(H(B_pk), H(A_pk), H(share ID), H(Nonce_A), H(Nonce_B))
  SecureChannelAuthMessage outmsg;
  setAuthMsg(me->keyhash(), peer->keyhash(), share->idhash(), nonce_a_hash,
             nonce_b_hash, &outmsg);
  
  if (!sendEncryptedMessage(sock, peer_auth_key, outmsg.SerializeAsString())) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't send at auth step 6.";
    return fail;
  }

  // 7. A: checks the same conditions as B did for C_1
  // (do nothing)

  // 8. A: stores H(Nonce_B) and the other values
  // (do nothing)

  // 9. A->B: C_3 = AE_{PA-key}(H(A_pk), H(B_pk), H(share ID), Nonce_A, H(Nonce_B))
  SecureChannelAuthMessage inmsg2;
  if (!processAuthMsg(instr, me, share, &inmsg2)) {
    DLOG(WARNING) << Q_FUNC_INFO << "auth message was in a bad format at auth step 9.";
    return fail;
  }

  std::string nonce_a = inmsg2.first_nonce();

  // 10. B: checks all the above conditions and checks that H(Nonce_A) == previously stored H(Nonce_A)
  if (crypto::hash(nonce_a) != nonce_a_hash) {
    DLOG(WARNING) << Q_FUNC_INFO << "nonce failed check at auth step 10.";
    return fail;
  }

  // 11. B->A: C_4 = AE_{PA-key}(H(B_pk), H(A_pk), H(share ID), Nonce_A, Nonce_B)
  SecureChannelAuthMessage outmsg2;
  setAuthMsg(me->keyhash(), peer->keyhash(), share->idhash(), nonce_a,
             nonce_b, &outmsg2);
 
  const std::string salt = outmsg2.SerializeAsString();
  const std::string &outmsg2_str = salt;

  if (!sendEncryptedMessage(sock, peer_auth_key, outmsg2_str)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't send at auth step 11.";
    return fail;
  }

  // 12. A: checks all the above conditions and checks that H(Nonce_B) == previously stored H(Nonce_B)
  // (do nothing)

  // 13. A: generates Diffie-Hellman key exchange parameters
  // (do nothing)

  // 14. A->B: AE_{PA-key}(PKE_{B_pk}(public_a))
  // receive the encrypted message
  if (!recvEncryptedMessage(sock, peer_auth_key, &instr)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't receive at auth step 14.";
    return fail;
  }
  
  // decrypt using our private key.
  std::string public_a_str;
  if (!crypto::pubkeyDecrypt(me->privkey(), instr, &public_a_str)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't decrypt at auth step 14.";
    return fail;
  }
  
  // 15. B: generates Diffie-Hellman key exchange parameters
  Botan::AutoSeeded_RNG rng;
  
  Botan::DL_Group shared_domain(crypto::kDHDomain);
  Botan::DH_PrivateKey private_b(rng, shared_domain);
  Botan::MemoryVector<Botan::byte> public_b = private_b.public_value();

  std::string public_b_str(reinterpret_cast<const char *>(public_b.begin()), public_b.size());

  std::string enc_public_b_str; // the public_a value encrypted under A's pubkey
  if (!crypto::pubkeyEncrypt(peer->pubkey(), public_b_str, &enc_public_b_str)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't encrypt at auth step 15.";
    return fail;
  }

  // 16. B->A: AE_{PA-key}(PKE_{A_pk}(public_b))
  if (!sendEncryptedMessage(sock, peer_auth_key, enc_public_b_str)) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't send at auth step 16.";
    return fail;
  }
  
  // 17. Generate session key from DH and salt
  Botan::MemoryVector<Botan::byte> public_a(reinterpret_cast<const Botan::byte *>(public_a_str.data()), public_a_str.size());
  Botan::PK_Key_Agreement ka(private_b, crypto::kDH_KDF);

  try {
    *session_key = ka.derive_key(crypto::kKeyLength, public_a, salt);
  } catch (Botan::Exception &e) {
    DLOG(ERROR) << Q_FUNC_INFO << "Botan exception: " << e.what();
    return fail;
  }

  return std::make_pair(peer, share);
}

}  // namespace octopod
